
#include "xdns-encoder.h"
#include "xdns-data.h"
#include "xdns-protocol.h"
#include <assert.h>

using namespace std;

XdnsError
XdnsEncoder::GenerateData(unsigned short id,
                          unsigned short qtype,
                          unsigned short qclass,
                          const string& name,
                          XdnsData& data)
{
    if (name.empty() ||
        (name.compare(0, 1, ".") == 0))
    {
        return XE_Inval;
    }

    data.Clear();

    /* The total data size should be saved in the first 2 bytes for
       TCP to send. The empty buffer is appended here, the data size
       will be set to these 2 bytes after all data composed.
    */
    unsigned char tcpDataSizeBuf[2] = {0};
    XdnsError error = data.AppendData(tcpDataSizeBuf,
                                      sizeof(tcpDataSizeBuf));
    if (error != XE_Ok)
    {
        data.Clear();
        return error;
    }

    // Compose the DNS header.
    unsigned char headerBuf[XDNS_HEADER_SIZE] = {0};
    XDNS_SET_HEADER_ID(headerBuf, id);
    XDNS_SET_HEADER_OPCODE(headerBuf, XDNS_OPCODE_QUERY);
    XDNS_SET_HEADER_RD(headerBuf, 1);
    XDNS_SET_HEADER_QDCOUNT(headerBuf, 1);
    error = data.AppendData(headerBuf, sizeof(headerBuf));
    if (error != XE_Ok)
    {
        data.Clear();
        return error;
    }

    /* Compose the QNAME. The query name is encoded as a list of
       labels, each label is represented by a one-byte size (max 63)
       followed by the text of the label. The list is terminated by
       a label of zero size which is considered as the root domain.
    */

    // The label size is saved in the first byte.
    unsigned char label[XDNS_MAX_LABEL_SIZE + 1] = {0};
    size_t offset = 1;

    string::const_iterator iter = name.begin();
    string::const_iterator end = name.end();
    for (; iter != end; ++iter)
    {
        if ('.' == (*iter))
        {
            if (offset > 1)
            {
                label[0] = offset - 1;
                error = data.AppendData(label, offset);
                if (error != XE_Ok)
                {
                    data.Clear();
                    return error;
                }

                offset = 1;
                continue;
            }

            data.Clear();
            return XE_Inval;
        }

        if ('\\' == (*iter) && ((iter + 1) != end))
        {
            continue;
        }

        label[offset] = static_cast<unsigned char>(*iter);
        ++offset;
        if (offset > sizeof(label))
        {
            data.Clear();
            return XE_Inval;
        }
    }

    // Append the last label.
    if (offset > 1)
    {
        label[0] = offset - 1;
        error = data.AppendData(label, offset);
        if (error != XE_Ok)
        {
            data.Clear();
            return error;
        }
    }

    // Append a null label at the end.
    unsigned char nullLabel = 0;
    error = data.AppendData(&nullLabel , 1);
    if (error != XE_Ok)
    {
        data.Clear();
        return error;
    }

    // Set the QTYPE and QCLASS.
    unsigned char qtypeBuf[4] = {0};
    XDNS_SET_QUESTION_TYPE(qtypeBuf, qtype);
    XDNS_SET_QUESTION_CLASS(qtypeBuf, qclass);
    error = data.AppendData(qtypeBuf, sizeof (qtypeBuf));
    if (error != XE_Ok)
    {
        data.Clear();
        return error;
    }

    /* Set the total data size (-2) to the first 2 bytes for TCP to
       send.
    */
    size_t totalDataSize = data.GetDataSize();
    assert(totalDataSize > 2);
    XDNS_SET_16BIT(tcpDataSizeBuf,
                   static_cast<unsigned short>(totalDataSize - 2));
    error = data.SetData(0, tcpDataSizeBuf, sizeof(tcpDataSizeBuf));
    if (error != XE_Ok)
    {
        data.Clear();
        return error;
    }

    return XE_Ok;
}

XdnsError
XdnsEncoder::ResetId(unsigned short id,
                     XdnsData& data)
{
    unsigned char idBuf[2] = {0};
    XDNS_SET_16BIT(idBuf, id);

    // The total data size is saved in the first 2 bytes.
    return data.SetData(2, idBuf, sizeof(idBuf));
}

